import tensorflow as tf


def image_gradient(x):
    """
    :param x: 5D tensor
    :return: 3 * 5D tensor
    """
    crop_1 = x[:, 0:-2, 1:-1, 1:-1, :]
    crop_2 = x[:, 1:-1, 0:-2, 1:-1, :]
    crop_3 = x[:, 1:-1, 1:-1, 0:-2, :]
    center = x[:, 1:-1, 1:-1, 1:-1, :]
    return center - crop_1, center - crop_2, center - crop_3


def decoder_loss(y_true, y_pred):
    mse_loss = tf.reduce_mean(tf.square(y_true - y_pred))
    g1_true, g2_true, g3_true = image_gradient(y_true)
    g1_pred, g2_pred, g3_pred = image_gradient(y_pred)
    gradient_loss = tf.reduce_mean(tf.square(g1_true - g1_pred)) + tf.reduce_mean(
        tf.square(g2_true - g2_pred)) + tf.reduce_mean(tf.square(g3_true - g3_pred))
    return mse_loss + 0.2 * gradient_loss
